from docx import Document
from docx.shared import Inches, Pt
from docx.enum.text import WD_ALIGN_PARAGRAPH
import io
import re
import jieba
import jieba.analyse
import numpy as np
from collections import Counter

jieba.setLogLevel(20)  # 静默jieba日志

def extract_keywords(text, topK=10):
    """
    提取文本关键词
    参数:
        text: 输入文本
        topK: 返回关键词数量
    返回:
        list: 关键词列表
    """
    # 移除HTML标签和特殊字符
    text = re.sub(r'<[^>]+>', '', text)
    text = re.sub(r'[^一-龥a-zA-Z0-9\s]', '', text)
    
    # 提取关键词
    keywords = jieba.analyse.extract_tags(text, topK=topK, allowPOS=('n', 'nr', 'ns', 'nt', 'nz'))
    
    # 如果jieba提取失败，回退到简单分词统计
    if not keywords:
        words = jieba.cut(text)
        filtered_words = [w for w in words if len(w) >= 2 and w.strip()]
        word_counts = Counter(filtered_words)
        keywords = [w for w, _ in word_counts.most_common(topK)]
    
    return keywords

def calculate_similarity(text1, text2):
    """
    计算两段文本的相似度，针对图片描述和段落内容进行优化
    参数:
        text1: 第一段文本（通常是段落内容）
        text2: 第二段文本（通常是图片描述）
    返回:
        float: 相似度分数 (0-1)
    """
    # 提取关键词，增加topK以更好地匹配前置处理的详细描述
    keywords1 = extract_keywords(text1, topK=20)  # 段落内容提取更多关键词
    keywords2 = extract_keywords(text2, topK=15)  # 图片描述关键词
    
    # 如果任一文本没有关键词，返回最低相似度
    if not keywords1 or not keywords2:
        return 0.0
    
    # 计算关键词交集
    common_keywords = set(keywords1) & set(keywords2)
    
    # 计算Jaccard相似度
    jaccard_similarity = len(common_keywords) / len(set(keywords1) | set(keywords2))
    
    # 考虑关键词顺序的相似度（简单的位置权重）
    ordered_similarity = 0
    for i, kw1 in enumerate(keywords1[:8]):  # 增加考虑的关键词数量
        if kw1 in keywords2:
            # 位置越靠前，权重越高
            pos_weight = (8 - i) / 8.0
            # 在第二个文本中的位置权重
            kw2_pos = keywords2.index(kw1)
            if kw2_pos < 8:  # 增加考虑的位置数量
                pos_weight2 = (8 - kw2_pos) / 8.0
                ordered_similarity += pos_weight * pos_weight2
    ordered_similarity = ordered_similarity / 8.0  # 归一化
    
    # 增加额外的文本相似度计算：TF-IDF匹配
    # 计算词频匹配度
    freq_match_score = 0
    text1_lower = text1.lower()
    text2_lower = text2.lower()
    
    # 计算共同词的出现频率
    for keyword in common_keywords:
        if keyword in text1_lower and keyword in text2_lower:
            freq_match_score += 0.1  # 为每个共同词增加基础分数
    
    # 限制最大值
    freq_match_score = min(freq_match_score, 0.3)
    
    # 综合相似度：60% Jaccard相似度 + 30% 顺序相似度 + 10% 词频匹配度
    similarity = 0.6 * jaccard_similarity + 0.3 * ordered_similarity + 0.1 * freq_match_score
    
    return similarity

def create_news_word_document(title, content, user_images=None, ai_image=None):
    """
    创建包含标题、内容和多张图片的Word文档，并将图片插入到最合适的位置
    参数:
        title: 新闻标题
        content: 新闻内容
        user_images: 用户上传的图片列表，每个元素是(image, description)元组，其中description是前置处理已获取的图片描述
        ai_image: AI生成的图片
    返回:
        bytes: Word文档的二进制数据
    """
    # 创建新文档
    doc = Document()
    # 设置文档字体
    style = doc.styles['Normal']
    font = style.font
    font.name = '宋体'
    font.size = Pt(12)
    
    # 添加标题
    title_para = doc.add_heading(title, level=0)
    title_para.alignment = WD_ALIGN_PARAGRAPH.CENTER
    title_run = title_para.runs[0]
    title_run.font.name = '黑体'
    title_run.font.bold = True
    title_run.font.size = Pt(18)
    
    # 如果有AI生成的图片，先放在标题和正文之间
    if ai_image:
        doc.add_paragraph()  # 空行
        insert_image_at_position(doc, ai_image)
        doc.add_paragraph()  # 空行分隔
    
    # 创建段落列表，去除空行并保留重要段落
    paragraphs = [p for p in content.split('\n') if p.strip()]
    
    # 如果有用户上传的图片，实现基于内容相关性的智能插入
    if user_images:
        # 直接使用前置处理提供的(image, description)格式
        processed_images = user_images
        
        # 提取整个文档的关键词，用于全局相关性分析
        document_keywords = extract_keywords(title + " " + content, topK=20)
        
        # 为每个段落提取关键词
        paragraphs_with_info = []
        for para_idx, para_text in enumerate(paragraphs):
            # 提取段落关键词
            para_keywords = extract_keywords(para_text, topK=8)
            # 计算段落长度（字符数）
            para_length = len(para_text)
            # 计算段落位置权重（开头和结尾的段落更重要）
            if para_idx == 0 or para_idx == len(paragraphs) - 1:
                position_weight = 0.9
            elif para_idx <= 2 or para_idx >= len(paragraphs) - 3:
                position_weight = 0.7
            else:
                position_weight = 0.5
            
            paragraphs_with_info.append({
                'text': para_text,
                'keywords': para_keywords,
                'length': para_length,
                'position_weight': position_weight,
                'idx': para_idx
            })
        
        # 为每个图片分配最相关的段落位置
        image_assignments = {}
        assigned_paragraphs = set()
        
        # 优先处理有明确描述的图片
        for img_idx, (image, description) in enumerate(processed_images):
            best_match_idx = -1
            best_match_score = 0.0
            
            # 对每个段落计算与图片的相关性
            for para_info in paragraphs_with_info:
                # 跳过已经分配了图片的段落
                if para_info['idx'] in assigned_paragraphs:
                    continue
                
                # 充分利用前置处理提供的图片描述计算相似度
                # 直接计算段落与完整图片描述的相似度
                similarity = calculate_similarity(para_info['text'], description)
                
                # 增强图片描述的权重，因为前置处理已确保描述质量和相关性
                if description and len(description) > 20:  # 前置处理限制在100字左右，这里优化权重计算
                    similarity *= 1.4  # 更长的描述通常包含更多相关信息
                elif description and len(description) > 10:
                    similarity *= 1.2
                
                # 结合文档级相关性，提高匹配准确性
                # 前置处理的描述已优化，使用较低阈值触发补充相关性
                if similarity < 0.1:  # 降低阈值，优先使用直接相似度
                    # 使用段落关键词与文档关键词的匹配度作为补充
                    common_keywords = set(para_info['keywords']) & set(document_keywords)
                    doc_match_score = len(common_keywords) / max(1, len(para_info['keywords']))
                    similarity = similarity * 0.7 + doc_match_score * 0.3  # 增加直接相似度权重
                
                # 应用位置权重
                weighted_score = similarity * para_info['position_weight']
                
                # 优先选择内容丰富的段落（长度适中的段落）
                if 80 <= para_info['length'] <= 400:  # 调整段落长度范围，使其更合理
                    weighted_score *= 1.3  # 增加权重提升效果
                
                # 更新最佳匹配
                if weighted_score > best_match_score:
                    best_match_score = weighted_score
                    best_match_idx = para_info['idx']
            
            # 分配图片到最佳匹配段落，或回退到均匀分布
            if best_match_idx >= 0 and best_match_score > 0.15:  # 调整阈值，使分配更准确
                image_assignments[best_match_idx] = (img_idx, image, description)
                assigned_paragraphs.add(best_match_idx)
            else:
                # 改进的回退策略：优先选择相关性次高且未分配的段落
                # 收集所有未分配段落的相似度分数
                available_scores = []
                for para_info in paragraphs_with_info:
                    if para_info['idx'] not in assigned_paragraphs:
                        # 重新计算相似度，充分利用前置处理的图片描述
                        sim = calculate_similarity(para_info['text'], description)
                        # 根据描述长度调整权重
                        if description and len(description) > 20:
                            sim *= 1.4
                        elif description and len(description) > 10:
                            sim *= 1.2
                        weighted_sim = sim * para_info['position_weight']
                        if 80 <= para_info['length'] <= 400:
                            weighted_sim *= 1.3
                        available_scores.append((weighted_sim, para_info['idx']))
                
                # 如果有可用段落，选择相关性最高的
                if available_scores:
                    # 按相似度降序排序
                    available_scores.sort(reverse=True)
                    # 选择第一个（相关性最高的）
                    _, best_available_idx = available_scores[0]
                    image_assignments[best_available_idx] = (img_idx, image, description)
                    assigned_paragraphs.add(best_available_idx)
        
        # 处理仍然未分配的图片，采用更智能的分配策略
        # 1. 找出所有未分配的图片和段落
        unassigned_img_indices = [idx for idx in range(len(processed_images)) 
                                 if idx not in [img[0] for img in image_assignments.values()]]
        unassigned_para_indices = [idx for idx in range(len(paragraphs)) 
                                  if idx not in image_assignments]
        
        # 2. 优先将未分配图片分配到未分配段落中最适合的位置
        for img_idx in unassigned_img_indices:
            if not unassigned_para_indices:
                break  # 没有未分配段落了
                
            image, description = processed_images[img_idx]
            best_para_idx = -1
            best_score = 0.0
            
            # 找到最佳匹配的未分配段落
            for para_idx in unassigned_para_indices:
                para_info = next(p for p in paragraphs_with_info if p['idx'] == para_idx)
                # 计算相似度，充分利用前置处理的图片描述
                sim = calculate_similarity(para_info['text'], description)
                # 根据描述长度调整权重
                if description and len(description) > 20:
                    sim *= 1.4
                elif description and len(description) > 10:
                    sim *= 1.2
                weighted_sim = sim * para_info['position_weight']
                if 80 <= para_info['length'] <= 400:
                    weighted_sim *= 1.3
                    
                if weighted_sim > best_score:
                    best_score = weighted_sim
                    best_para_idx = para_idx
            
            # 如果找到合适的段落，进行分配
            if best_para_idx >= 0:
                image_assignments[best_para_idx] = (img_idx, image, description)
                unassigned_para_indices.remove(best_para_idx)
        
        # 3. 如果还有未分配的图片，尝试添加到开头或结尾
        remaining_img_indices = [idx for idx in unassigned_img_indices 
                                if idx not in [img[0] for img in image_assignments.values()]]
        for img_idx in remaining_img_indices:
            image, description = processed_images[img_idx]
            # 尝试添加到开头
            if 0 not in image_assignments:
                image_assignments[0] = (img_idx, image, description)
            # 或添加到结尾
            elif len(paragraphs) - 1 not in image_assignments:
                image_assignments[len(paragraphs) - 1] = (img_idx, image, description)
        
        # 构建最终文档：按照段落顺序添加内容和图片
        for para_idx, para_text in enumerate(paragraphs):
            # 添加段落内容
            add_formatted_paragraph(doc, para_text)
            
            # 检查当前段落是否分配了图片
            if para_idx in image_assignments:
                img_idx, image, description = image_assignments[para_idx]
                doc.add_paragraph()  # 空行
                insert_image_at_position(doc, image)
        
        # 处理未分配的图片，添加到文档末尾
        used_image_indices = [img[0] for img in image_assignments.values()]
        for img_idx, (image, description) in enumerate(processed_images):
            if img_idx not in used_image_indices:
                doc.add_paragraph()  # 空行
                insert_image_at_position(doc, image)
    else:
        # 如果没有图片，直接添加所有内容
        for para_text in paragraphs:
            add_formatted_paragraph(doc, para_text)
    
    # 保存文档到字节流
    doc_byte_arr = io.BytesIO()
    doc.save(doc_byte_arr)
    doc_byte_arr.seek(0)
    
    return doc_byte_arr.getvalue()

def add_formatted_paragraph(doc, text):
    """
    添加格式化的段落到文档
    参数:
        doc: Word文档对象
        text: 段落文本
    """
    # 检查是否是标题格式
    if text.startswith('# '):
        # 小标题黑体字(不加粗)
        heading_text = text[2:].strip()
        para = doc.add_heading(heading_text, level=2)
        run = para.runs[0]
        run.font.name = '黑体'
        run.font.size = Pt(14)
        return
    
    # 普通段落
    para = doc.add_paragraph(text)
    run = para.runs[0]
    run.font.name = '宋体'
    run.font.size = Pt(12)
    
    # 设置段落格式
    para.paragraph_format.space_after = Pt(6)
    para.paragraph_format.space_before = Pt(6)
    para.paragraph_format.line_spacing = 1.15

def insert_image_at_position(doc, image):
    """
    在文档指定位置插入图片
    参数:
        doc: Word文档对象
        image: PIL Image对象
    """
    # 将PIL Image转换为BytesIO对象
    img_byte_arr = io.BytesIO()
    image.save(img_byte_arr, format='PNG')
    img_byte_arr.seek(0)
    
    # 调整图片大小，确保不会太大
    img_width = Inches(6)
    # 插入图片
    doc.add_picture(img_byte_arr, width=img_width)
    # 设置图片居中对齐
    doc.paragraphs[-1].alignment = WD_ALIGN_PARAGRAPH.CENTER